import pyaudio
import json
from PyQt5.QtCore import QThread
from vosk import Model, KaldiRecognizer
from utils import findInternalRecordingDevice

# 设置音频流参数
RATE = 44100  # 采样率
CHUNK = 2048  # 每次读取的音频块大小

class RecognitionThread(QThread):
    def __init__(self, window):
        super().__init__()
        self.window = window

    # 识别线程的 run 方法
    def run(self):
        p = pyaudio.PyAudio()
        device_index = findInternalRecordingDevice(p)
        if device_index == -1:
            print("没有找到内录设备，程序终止。")
            return

        model = Model("vosk-model-cn-0.22")
        rec = KaldiRecognizer(model, RATE)

        stream = p.open(input_device_index=device_index,
                        format=pyaudio.paInt16,
                        channels=1,
                        rate=RATE,
                        input=True,
                        frames_per_buffer=CHUNK)

        try:
            while True:
                data = stream.read(CHUNK, exception_on_overflow=False)

                if rec.AcceptWaveform(data):
                    result = rec.Result()
                    result_json = json.loads(result)

                    final_text = result_json.get('text', '').replace(" ", "")
                    if final_text:
                        # 将完整句子追加到历史中
                        self.window.history.append(final_text)
                        self.window.transcription = ""
                        print(f"Stored full sentence: {final_text}")

                else:
                    partial_result = rec.PartialResult()
                    partial_text = json.loads(partial_result).get('partial', '').replace(" ", "")
                    # 实时覆盖当前的字幕显示，但保留历史
                    self.window.transcription = partial_text
                    print(f"Partial result: {partial_text}")

        except KeyboardInterrupt:
            print("录音结束")

        stream.stop_stream()
        stream.close()
        p.terminate()
